import oneflow as torch
import oneflow.nn as nn
import oneflow.nn.functional as F
from oasr.data import PAD

class LabelSmoothingLoss(nn.Module):
    """Label-smoothing loss

    :param int size: the number of class
    :param int padding_idx: ignored class id
    :param float smoothing: smoothing rate (0.0 means the conventional CE)
    :param bool normalize_length: normalize loss by sequence length if True
    :param torch.nn.Module criterion: loss function to be smoothed
    """

    def __init__(self, size, smoothing, padding_idx=PAD, normalize_length=True,
                 criterion=nn.KLDivLoss(reduction='none')):
        super(LabelSmoothingLoss, self).__init__()
        self.criterion = criterion
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None
        self.normalize_length = normalize_length

    def forward(self, x, target):
        """Compute loss between x and target

        :param torch.Tensor x: prediction (batch, seqlen, class)
        :param torch.Tensor target: target signal masked with self.padding_id (batch, seqlen)
        :return: scalar float value
        :rtype torch.Tensor
        """
        assert x.size(2) == self.size
        batch_size = x.size(0)
        x = x.view(-1, self.size)
        target = target.reshape(-1)

        with torch.no_grad():
            true_dist = x.clone()
            true_dist.fill_(self.smoothing / (self.size - 1))
            ignore = target == self.padding_idx  # (B,)
            total = len(target) - ignore.sum().item()
            target = target.masked_fill(ignore, 0)  # avoid -1 index
            true_dist.scatter_(1, target.unsqueeze(1), self.confidence)

        kl = self.criterion(torch.log_softmax(x, dim=1), true_dist)
        denom = total if self.normalize_length else batch_size

        return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom


class MaskedLabelSmoothingLoss(nn.Module):
    def __init__(self, size, smoothing=0.1, padding_idx=PAD, normalize_length=True):
        super().__init__()

        self.padding_idx = padding_idx
        self.smoothing = smoothing
        self.size = size
        self.normalize_length = normalize_length

    def forward(self, logits, target, mask=None):
        """LabelSmoothing Function with Mask

        Args:
            logits ([tensor]): logits with shape [batch, length, vocab_size]
            target ([tensor]): target with shape [batch, length]
            mask ([tensor], optional): mask tensor (bool) with shape [batch, length]
        """
        assert logits.dim() == 3 and logits.size(-1) == self.size
        
        pad_mask = target == self.padding_idx
        if mask is not None:
            mask = (pad_mask.int() + mask.int()) > 0
        else:
            mask = pad_mask

        logits = logits.reshape(-1, self.size)
        with torch.no_grad():
            confidence = logits.clone()
            confidence.fill_(self.smoothing / (self.size - 1))
            confidence.scatter_(1, target.reshape(-1).unsqueeze(1), 1 - self.smoothing)

        loss = torch.sum(F.kl_div(F.log_softmax(logits, dim=-1), confidence, reduction='none'), dim=-1)
        total = torch.sum(~mask)
        denom = total if self.normalize_length else logits.size(0)
        loss = torch.sum(loss.masked_fill_(mask.reshape(-1), 0.0)) / denom

        return loss


class KLDivergenceLoss(nn.Module):
    def __init__(self, normalize_length=True) -> None:
        super().__init__()

        self.normalize_length = normalize_length

    def forward(self, xlogits, ylogits, mask):
        """Compute KLDivergence Loss with Mask

        Args:
            xlogits ([tensor]): xlogits tensor with shape [b, t, v]
            ylogits ([tensor]): ylogits tensor with shape [b, t, v]
            mask ([tensor]): Bool Mask Tensor with shape [b, t].
        """

        klloss = F.kl_div(F.log_softmax(xlogits, dim=-1), F.softmax(ylogits, dim=-1), reduction='none')
        klloss = torch.sum(klloss, dim=-1)
        klloss.masked_fill_(mask, 0.0)

        demon = torch.sum(~mask) if self.normalize_length else xlogits.size(0)
        return torch.sum(klloss) / demon
        

class CumulativeCrossEntropyLoss(nn.Module):
    def __init__(self, smoothing=0.0, padding_idx=PAD):
        super().__init__()

        self.smoothing = smoothing
        self.padding_idx = padding_idx

    def forward(self, log_probs, target):
        """CumulativeCrossEntropyLoss for transducer

        Args:
            log_probs ([tensor]): logits with shape [b, t, u, v]
            target ([tensor]): target with shape [b, t]
        """
        assert log_probs.dim() == 4 and target.dim() == 2 and target.size(1) == log_probs.size(2)
        _, t, _, v = log_probs.size()
        target = target.unsqueeze(1).repeat([1, t, 1])
        mask = target == self.padding_idx

        log_probs = log_probs.reshape(-1, v)
        with torch.no_grad():
            confidence = log_probs.clone()
            confidence.fill_(self.smoothing / (v - 1))
            confidence.scatter_(1, target.reshape(-1).unsqueeze(1), 1 - self.smoothing)

        loss = torch.sum(F.kl_div(log_probs, confidence, reduction='none'), dim=-1)
        total = torch.sum(~mask)
        loss = torch.sum(loss.masked_fill_(mask.reshape(-1), 0.0)) / total

        return loss